
#include "HookProcManagerImpl_Win.h"





#include SD_TL_PLATFORM_ARH_INC( .,HookProc )
#include "tools/impl/win/WinExpection.h"
#include "tools/locker.h"
#include "tools/ThreadEnumerator.h"
#include SD_TL_PLATFORM_IMPL_INC( tools,HandleWraper )






namespace sd
{
namespace hp
{
namespace impl
{

using namespace sd::tools::buffer;
using namespace sd::tools::win;




struct jumper
{
private:
	StaticBuffer<100> data_;

	PtrBuffer injector_;
	PtrBuffer caller_;
	PtrBuffer repeater_;
	PtrBuffer resuamer_;

		
public:
	jumper():
	  injector_(0,0), caller_(0,0), repeater_(0,0), resuamer_(0,0)	{}

	const PtrBuffer& getInject()const
	{ 
		return injector_;  
	}
	
	void setOldData( OldFunctionData* olddata )
	{
		data_.clear();
		resuamer_.reset( pushInBuffer( olddata->createResumer( data_.end() ) ) );
		caller_.reset( pushInBuffer( olddata->createCaller( data_.end(), resuamer_.begin() ) ) );
		injector_.reset( pushInBuffer( olddata->createInjector( caller_.begin() ) ) );
		repeater_.reset( pushInBuffer( olddata->createRepeater(  data_.end() ) ) );
	}

	void dropOldData()
	{
		data_.clear();
		resuamer_.clear();
		caller_.clear();
		injector_.clear();
		repeater_.clear();
	}

	const Buffer& getData()const{ return data_; }

	const void* getOriginalCaller()
	{
		return repeater_.begin();
	}

private:
	PtrBuffer pushInBuffer( const Buffer& buffer )
	{
		const void* old = data_.end();
		data_ += buffer;
		return PtrBuffer( old, buffer.size() );
	}
};






class ExceptionNoStandartFunction: public std::exception
{
public:
	~ExceptionNoStandartFunction()throw(){}
	ExceptionNoStandartFunction( const ExceptionNoStandartFunction& ref )throw(): name_( ref.name_ ){}
	
	
	ExceptionNoStandartFunction( const std::string& name )throw(): name_(name){}

	virtual const char* what()const throw()
	{
		return std::string( "Call function without standart prolog: "+ name_ ).c_str();
	}
private:
	ExceptionNoStandartFunction()throw(){}
	std::string name_;
};





LibraryImpl::LibraryImpl( const std::string& name ): name_( name )
{
	lib_ = LoadLibraryA( name.c_str() );
	if ( !lib_ )
		throw WinImplException( std::string("Library ") + name + " not load" );
}



UntypedFunction LibraryImpl::getFunction( const std::string& name )const
{
    FARPROC proc = GetProcAddress( lib_, name.c_str() );
    if ( !proc )
        throw WinImplException( "Function '" + name + "' not found in module '" + getName() + "'" );
	return UntypedFunction( *this, name, proc );
}


LibraryImpl::~LibraryImpl()
{
	FreeLibrary( lib_ );
}





HookImplementation::HookImplementation( const UntypedFunction& function, Callback callback, void* data ): 
	function_( function ), data_( data ), callback_( callback ), active_( false ), oldData_( 0 ), jumper_( new jumper ),lockid(0),counter(0)
{} 

void HookImplementation::setOldData( OldFunctionData* data )
{
	if ( data )
		jumper_->setOldData( data );
	else
		jumper_->dropOldData();

	oldData_ = data;
}


void HookImplementation::setActive( bool active )
{
	if( isActive() != active )
	{
		if( active )
			HookedProcessImplementation::getInstance().SetHook( *this );
		else
			HookedProcessImplementation::getInstance().removeHook( *this );
		active_ = active;
	}
}


void ProcessLocker::lock()
{
	assert( !locked_ );
	if(locked_)return;

	const int MAX_LOCK_TRY = 10;
	int i;
	for( i=0; i < MAX_LOCK_TRY && tryLock(); i++ );
	locked_ = true;
	if( i >= MAX_LOCK_TRY )
	{
		unlock();
		throw WinImplException( "Process lock not completed" );
	}
}

int ProcessLocker::tryLock()
{
	using namespace tools;
	using namespace thread_enumerator;

	ProcessThreadEnumerator process( (HHANDLE( GetCurrentProcess() ) ) );

	int result = 0;

	thread::Thread current;

	for( ProcessThreadEnumerator::iterator it = process.begin(); it != process.end(); ++it )
		if( *it != current )
		{
			if( syspended_threads.find( *it ) == -1 )
			{
				if(! (*it).isSuspended())
				{
					(*it).suspend();
					syspended_threads.addThread(*it);
					result++;
				}
			}
		}
	return result;
}

void ProcessLocker::unlock()
{
	assert( locked_ );
	if(!locked_)return;

	syspended_threads.resume();
	syspended_threads.clear();
	locked_=false;
}


bool ProcessLocker::checkMemoryRegionExec( const void* start_region, const void* end_region )const
{
	for(size_t i = 0; i < syspended_threads.count(); ++i)
	{
		const void* EIP = syspended_threads.get(i).getCurrentEIP();
		if(EIP >= start_region && EIP < end_region)
			return true;
	}
	return false;

}


HookedProcessImplementation::HookedProcessImplementation()
{}

HookedProcessImplementation::~HookedProcessImplementation()
{
	tools::locker<HookedProcessImplementation> lock(*this);

	for( int index = hooks_.size()-1; index >=0 ; --index )
		removeHook( index );
	hooks_.clear();
}

HookImplementationRef& HookedProcessImplementation::CreateHook( const LibraryImpl& lib,  const std::string& name, Callback callback, void* param )
{
	UntypedFunction func = lib.getFunction( name );
	size_t hookIndex = AddHook( HookImplementationRef( new HookImplementation( func, callback, param ) ) );
	return hooks_[ hookIndex ];	
}

HookImplementation::OriginalFunctionCallerPtr HookImplementation::getCallerOriginalPtr()const
{
	assert( isActive() );
    return (OriginalFunctionCallerPtr)jumper_->getOriginalCaller();
}





int HookedProcessImplementation::AddHook( const HookImplementationRef& ref )
{
	hooks_.push_back( ref );
	return hooks_.size() - 1;
}

void HookedProcessImplementation::SetHook( size_t index )
{
    assert( index >=0 && index < hooks_.size() );
    SetHook( *hooks_[ index ] );
}



void HookedProcessImplementation::SetHook( HookImplementation& hook )
{
    if ( hook.isActive() )
        return;
        
    OldFunctionData* oldData = 0;

	if ( OldFunctionDataApiProlog::checkData( hook.getFunction() ) )
        oldData = new OldFunctionDataApiProlog( &hook, hook.getFunction() );
	else if ( OldFunctionDataHook::checkData( hook.getFunction() ) )
		oldData = new OldFunctionDataHook( &hook, hook.getFunction() );

    if ( !oldData )
        throw ExceptionNoStandartFunction( hook.getFunction().getName() );

    hook.setOldData( oldData );
	const jumper& jump = hook.getJumper();
	const PtrBuffer& injector = jump.getInject();

	setMemoryFlag( jump.getData(), MemoryLocker::toWindows( LM_ALL ) );

	MemoryLocker locker( this, PtrBuffer( oldData->getPlacePointer(), injector.size() ), LM_ALL );
    writeMemory( locker, injector );
}


void HookedProcessImplementation::removeHook( size_t index )
{
    assert( index >=0 && index < hooks_.size() );
    removeHook( *hooks_[ index ] );
}


HookException::HookException( const HookImplementation& impl, const std::string& error ): impl_( impl )
{
	emsg_ = impl_.getFunction().getLibrary().getName() +"."+impl_.getFunction().getName()+":"+error;
}


void HookedProcessImplementation::removeHook( HookImplementation& hook )
{
    if ( !hook.isActive() )
        return;

	const PtrBuffer& injector = hook.getJumper().getInject();
	OldFunctionData* oldData =  hook.getOldData();
	
	MemoryLocker locker( this, PtrBuffer( oldData->getPlacePointer(), injector.size() ), LM_READ );

	if( injector != locker )
			throw HookException( hook, "Removing hook error, maybe hook is overrided" );

	writeMemory( locker, oldData->getOldData() );
	hook.setOldData( 0 );
}




void HookedProcessImplementation::writeMemory( MemoryLocker& lock, const Buffer& buffer )
{
	lock.setFlag( LM_WRITE );
	memcpy( const_cast<void*>(lock.begin()), buffer.begin(), buffer.size() > lock.size() ? lock.size() : buffer.size() );
}


bool HookedProcessImplementation::checkLockedThreadsActiveCode(const void* ptr)const
{

	const void *begin = ptr;
	const void *end = (const char*)ptr + sizeof(DefaultApiProlog) + sizeof(void*) + 1;
	return checkMemoryRegionExec(begin, end);

}




DWORD HookedProcessImplementation::setMemoryFlag( const Buffer& buffer, DWORD protect )
{
	DWORD oldProtect;
	if ( !VirtualProtect( const_cast<void*>(buffer.begin()), buffer.size(), protect, &oldProtect ) )
		throw WinImplException( "Call VirtualProtect fail" );
	return oldProtect;
}


}//namespace impl
}//namespace hp
}//namespace sd
